from __future__ import division

import numpy as np
from matplotlib import pyplot as plt

# Load the mandrill image as an NxNx3 array. Values range from 0.0 to 255.0.
mandrill = plt.imread('mandrill.png')[:,:,:3].astype(float)
N = int(mandrill.shape[0])

M = 2
k = 64

# Store each MxM block of the image as a row vector of X
X = np.zeros((N**2//M**2, 3*M**2))
for i in range(N//M):
    for j in range(N//M):
        X[i*N//M+j,:] = mandrill[i*M:(i+1)*M,j*M:(j+1)*M,:].reshape(3*M**2)

'''
TODO: Implement k-means and cluster the rows of X, then reconstruct the compressed image using the cluster center for each block, as specified in the homework description.
'''
def kmeans(X, k, max_iters=50, tol=1e-4):
    n, d = X.shape
    rng = np.random.default_rng(seed=0)
    centers = X[rng.choice(n, k, replace=False)]
    prev_obj = float('inf')
    obj_history = []

    for it in range(max_iters):
        # Assignment Step
        dists = np.linalg.norm(X[:, None, :] - centers[None, :, :], axis=2)
        labels = np.argmin(dists, axis=1)

        # Update Step
        new_centers = np.array([X[labels == i].mean(axis=0) if np.any(labels == i) else centers[i] for i in range(k)])

        # Objective function
        obj = np.sum((X - new_centers[labels]) ** 2)
        obj_history.append(obj)

        # Check convergence
        if abs(prev_obj - obj) < tol:
            break
        prev_obj = obj
        centers = new_centers

    return centers, labels, obj_history

# TODO: k-means
centers, labels, obj_history = kmeans(X, k)

compressed_X = centers[labels]
compressed_img = np.zeros_like(mandrill)

for i in range(N // M):
    for j in range(N // M):
        idx = i * (N // M) + j
        block = compressed_X[idx].reshape(M, M, 3)
        compressed_img[i*M:(i+1)*M, j*M:(j+1)*M, :] = block

# objective function
plt.figure()
plt.plot(obj_history)
plt.title("K-means Objective vs Iteration")
plt.xlabel("Iteration")
plt.ylabel("Objective Function Value")
plt.grid(True)
plt.savefig("objective_plot.png")

# original and compressed image
plt.figure()
plt.subplot(1, 2, 1)
plt.title("Original Image")
plt.imshow(mandrill.astype(np.uint8))

plt.subplot(1, 2, 2)
plt.title("Compressed Image")
plt.imshow(compressed_img.astype(np.uint8))
plt.savefig("comparison.png")

# plot difference
diff_img = (mandrill - compressed_img + 128).clip(0, 255).astype(np.uint8)
plt.figure()
plt.title("Difference Image (+128 Gray)")
plt.imshow(diff_img)
plt.savefig("difference.png")

# RMAE
RMAE = np.mean(np.abs(mandrill - compressed_img)) / 255
print("Relative Mean Absolute Error:", RMAE)